/**
 * 
 */
package info.niwota.webagent.util;

import java.util.Collection;
import java.util.Enumeration;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * Supports blocking for get operation
 * 
 * @author qiangli
 * @param <K>
 * @param <V>
 *
 */
public class BlockingMap<K, V> implements Map<K, V> {
	
	private class LatchedObject<T extends V> {
		public V object;
		public CountDownLatch latch;
	}
	
	private ConcurrentHashMap<K, LatchedObject<V>> map = null;
	
	public BlockingMap() {
		map = new ConcurrentHashMap<K, LatchedObject<V>>();
	}

	public void clear() {
		map.clear();
	}

	public boolean contains(Object value) {
		return map.contains(value);
	}

	public boolean containsKey(Object key) {
		return map.containsKey(key);
	}

	public boolean containsValue(Object value) {
		return map.containsValue(value);
	}

	public Enumeration<V> elements() {
		throw new UnsupportedOperationException();
	}

	public Set<java.util.Map.Entry<K, V>> entrySet() {
		throw new UnsupportedOperationException();
	}

	public V peek(Object key) {
		LatchedObject<V> o = map.get(key);
		if (o != null) {
			return o.object;
		}
		return null;
	}
	
	public void release(Object key) {
		LatchedObject<V> o = map.get(key);
		if (o != null) {
			o.latch.countDown();
		}
	}
	
	/**
	 * Block until available
	 */
	public V get(Object key) {
		return get(key, Long.MAX_VALUE);
	}
	
	/**
	 * Block until available or time out
	 * @param unit 
	 * @param timeout 
	 */
	public V get(Object key,  long timeout) {
		LatchedObject<V> o = map.get(key);
		if (o != null) {
			try {
				o.latch.await(timeout, TimeUnit.SECONDS);
			} catch (Exception e) {
				return null;
			} finally {
				o.latch.countDown();
				map.remove(key);
			}
			return o.object;
		}
		return null;
	}

	public int hashCode() {
		return map.hashCode();
	}

	public boolean isEmpty() {
		return map.isEmpty();
	}

	public Set<K> keySet() {
		return map.keySet();
	}

	public Enumeration<K> keys() {
		return map.keys();
	}

	public V put(K key, V value) {
		LatchedObject<V> o = new LatchedObject<V>();
		o.object = value;
		map.put(key, o);
		return value;
	}

	public void putAll(Map<? extends K, ? extends V> m) {
		throw new UnsupportedOperationException();
	}

	public V putIfAbsent(K key, V value) {
		LatchedObject<V> o = map.get(key);
		if (o == null) {
			LatchedObject<V> lvalue = new LatchedObject<V>();
			lvalue.object = value;
			map.putIfAbsent(key, lvalue);
			return value;
		}
		return o.object;
	}

	public boolean remove(Object key, Object value) {
		return map.remove(key, value);
	}

	public V remove(Object key) {
		LatchedObject<V> o = map.remove(key);
		return (o == null ? null : o.object);
	}

	public boolean replace(K key, V oldValue, V newValue) {
		LatchedObject<V> o = map.get(key);
		if (o != null && o.object != null && o.object.equals(oldValue)) {
				o.object = newValue;
				return true;
		}
		return false;
	}

	public V replace(K key, V value) {
		LatchedObject<V> o = map.get(key);
		if (o != null ) {
			o.object = value;
			return o.object;
		}
		return null;
	}

	public int size() {
		return map.size();
	}

	public Collection<V> values() {
		throw new UnsupportedOperationException();
	}
}
